#include "binary.h"

#include <cstring>
#include <functional>

#include "tinyndarray.h"

namespace tinyndarray {

// 通用的二元操作函数
void ApplyBinaryOp(NdArray& ret, const NdArray& lhs, const NdArray& rhs, BinaryElementOp op) {
    if (lhs.shape() == rhs.shape()) {
        ApplyOpSimple(ret, lhs, rhs, op);
    } else {
        Shape ret_shape = CheckBroadcastable(lhs.shape(), rhs.shape());
        ApplyOpBroadcast(ret, lhs, rhs, 0, 1, op);
    }
}

NdArray ApplyBinaryOp(const NdArray& lhs, const NdArray& rhs, BinaryElementOp op) {
    if (lhs.shape() == rhs.shape()) {
        NdArray ret(lhs.dtype(), lhs.shape());
        ApplyOpSimple(ret, lhs, rhs, op);
        return ret;
    } else {
        Shape ret_shape = CheckBroadcastable(lhs.shape(), rhs.shape());
        NdArray ret(lhs.dtype(), ret_shape);
        ApplyOpBroadcast(ret, lhs, rhs, 0, 1, op);
        return ret;
    }
}

#define BINARY_OP(NAME, OPR)                                                     \
    template <typename T>                                                        \
    void NAME##_op(void* out, const void* a, const void* b) {                    \
        T* out_typed = static_cast<T*>(out);                                     \
        const T* a_typed = static_cast<const T*>(a);                             \
        const T* b_typed = static_cast<const T*>(b);                             \
        *out_typed = *a_typed OPR *b_typed;                                      \
    }                                                                            \
                                                                                 \
    BinaryElementOp get_##NAME##_op(Dtype dtype) {                               \
        switch (dtype.tag) {                                                     \
            case FP32: {                                                         \
                return NAME##_op<float>;                                         \
            }                                                                    \
            case INT32: {                                                        \
                return NAME##_op<int>;                                           \
            }                                                                    \
            case FP64: {                                                         \
                return NAME##_op<double>;                                        \
            }                                                                    \
            default: throw std::runtime_error("Unsupported data type for name"); \
        }                                                                        \
        return nullptr;                                                          \
    }

BINARY_OP(add, +)
// 算术操作实现 - 数组与数组
// 加法操作的实现
NdArray Add(const NdArray& lhs, const NdArray& rhs) {
    BinaryElementOp typed_op = get_add_op(lhs.dtype());
    return ApplyBinaryOp(lhs, rhs, typed_op);
}

template <typename T>
NdArray Add(const NdArray& lhs, const T& rhs) {
    BinaryElementOp typed_op = get_add_op(lhs.dtype());
    return ApplyBinaryOp(lhs, rhs, typed_op);
}

// 算术运算符实现
NdArray operator+(const NdArray& lhs, const NdArray& rhs) {
    return Add(lhs, rhs);
}

template <typename T>
NdArray operator+(const NdArray& lhs, const T& rhs) {
    return Add(lhs, rhs);
}

template NdArray operator+(const NdArray& lhs, const float& rhs);

BINARY_OP(sub, -)

NdArray Subtract(const NdArray& lhs, const NdArray& rhs) {
    BinaryElementOp typed_op = get_sub_op(lhs.dtype());

    return ApplyBinaryOp(lhs, rhs, typed_op);
}

NdArray operator-(const NdArray& lhs, const NdArray& rhs) {
    return Subtract(lhs, rhs);
}

BINARY_OP(mul, *)

NdArray Multiply(const NdArray& lhs, const NdArray& rhs) {
    BinaryElementOp typed_op = get_mul_op(lhs.dtype());

    return ApplyBinaryOp(lhs, rhs, typed_op);
}

template <typename T>
NdArray Multiply(const NdArray& lhs, const T& rhs) {
    BinaryElementOp typed_op = get_mul_op(lhs.dtype());
    return ApplyBinaryOp(lhs, rhs, typed_op);
}

NdArray operator*(const NdArray& lhs, const NdArray& rhs) {
    return Multiply(lhs, rhs);
}

template <typename T>
NdArray operator*(const NdArray& lhs, const T& rhs) {
    return Multiply(lhs, rhs);
}

template NdArray operator*(const NdArray& lhs, const float& rhs);

BINARY_OP(div, /)

NdArray operator/(const NdArray& lhs, const NdArray& rhs) {
    BinaryElementOp typed_op = get_div_op(lhs.dtype());

    return ApplyBinaryOp(lhs, rhs, typed_op);
}

template <typename T>
NdArray Div(const NdArray& lhs, const T& rhs) {
    BinaryElementOp typed_op = get_div_op(lhs.dtype());
    return ApplyBinaryOp(lhs, rhs, typed_op);
}

template <typename T>
NdArray operator/(const NdArray& lhs, const T& rhs) {
    return Div(lhs, rhs);
}

template NdArray operator/(const NdArray& lhs, const float& rhs);

#define BINARY_CMP_OP(NAME, OPR)                                                 \
    template <typename T>                                                        \
    void NAME##_op(void* out, const void* a, const void* b) {                    \
        bool* out_typed = static_cast<bool*>(out);                               \
        const T* a_typed = static_cast<const T*>(a);                             \
        const T* b_typed = static_cast<const T*>(b);                             \
        *out_typed = *a_typed OPR *b_typed;                                      \
    }                                                                            \
                                                                                 \
    BinaryElementOp get_##NAME##_op(Dtype dtype) {                               \
        switch (dtype.tag) {                                                     \
            case FP32: {                                                         \
                return NAME##_op<float>;                                         \
            }                                                                    \
            case INT32: {                                                        \
                return NAME##_op<int>;                                           \
            }                                                                    \
            case FP64: {                                                         \
                return NAME##_op<double>;                                        \
            }                                                                    \
            default: throw std::runtime_error("Unsupported data type for name"); \
        }                                                                        \
        return nullptr;                                                          \
    }

BINARY_CMP_OP(eq, ==)

NdArray Equal(const NdArray& lhs, const NdArray& rhs) {
    BinaryElementOp typed_op = get_eq_op(lhs.dtype());
    NdArray ret(BOOL, lhs.shape());
    ApplyOpSimple(ret, lhs, rhs, typed_op, false);
    return ret;
}

NdArray operator==(const NdArray& lhs, const NdArray& rhs) {
    return Equal(lhs, rhs);
}

BINARY_CMP_OP(less, <)

NdArray Less(const NdArray& lhs, const NdArray& rhs) {
    BinaryElementOp typed_op = get_less_op(lhs.dtype());
    NdArray ret(BOOL, lhs.shape());
    ApplyOpSimple(ret, lhs, rhs, typed_op, false);
    return ret;
}

NdArray operator<(const NdArray& lhs, const NdArray& rhs) {
    return Less(lhs, rhs);
}
}  // namespace tinyndarray